﻿#include <cmath>
#include <QCryptographicHash>
#include "merkle.h"

const int DEPTH = 9;
const int ARITY = 8;

__device__ BitArray sha3_256_bool(const BitArray &vec_bool)
{
    CuByteArray vec_bytes = bits_to_bytes_le(vec_bool);

    QByteArray bytes(vec_bytes.data(),vec_bytes.size());
    QByteArray qret = QCryptographicHash::hash(bytes, QCryptographicHash::Sha3_256);

    CuByteArray ret = toBytes(qret.data(),qret.size());
    return bytes_to_bits_le(ret);
}

class LeaveHash
{
public:
    __device__ BitArray hash_leaf(const BitArray &leaf)
    {
        BitArray input;

        // Prepend the leaf with a `false` bit.
        Boolean b;
        b.value = false;
        input.push_back(b);
        input.append(leaf);

        // Hash the input.
        BitArray output = sha3_256_bool(input);
        // Read the first VARIANT bits. 
        return output.mid(0,256);
    }
};

struct PathHash
{   
    __device__ BitArray hash_children(const CuVector<BitArray>& children)
    {
        //qDebug().noquote() << "children input:" << bool_vec_to_hex_string(children);

        BitArray input;
        // Prepend the nodes with a `true` bit.
        Boolean b;
        b.value = true;
        input.push_back(b);
        for (int i = 0; i < children.size(); i++)
        {
            input.append(children[i]);
        }

        // Hash the input.
        BitArray output = sha3_256_bool(input);
        //qDebug() << "output:" << to_hex_string(bool_to_bytes_le(output));

        return output.mid(0,256);
    }

    __device__ BitArray hash_empty() 
    {
        Boolean b = Boolean::constant(false);
        CuVector<BitArray> children;
        for (int i = 0; i < ARITY; i++)
        {
            children.push_back(BitArray(256, b));
        }
        return hash_children(children);
    }
};



// 检查并返回大于等于base的以n为底数的最小幂次，类似Rust中的checked_next_power_of_n函数
__device__ int checked_next_power_of_n(int base, int n) {
    if (n <= 1) {
        return base;
    }
    int value = 1;
    while (value < base) {
        value *= n;
    }
    return value;
}

// 计算树的深度，模拟Rust中的tree_depth函数，这里使用double类型来模拟浮点数计算
__device__ int calc_tree_depth(int tree_size)
{
    double tree_depth_float = std::log(static_cast<double>(tree_size)) / std::log(static_cast<double>(ARITY));
    int tree_depth = static_cast<int>(std::floor(tree_depth_float));
    if (tree_depth > DEPTH) {
        assert(0);
    }
    return tree_depth;
}

struct Range
{
    int r0;
    int r1;
};

// 返回给定索引节点的子节点索引范围，模拟Rust中的child_indexes函数
__device__ Range child_indexes(int index) {
    int start = index * ARITY + 1;
    return {start, start + ARITY};
}

// 计算给定索引节点的父节点索引，模拟Rust中的parent函数
__device__ int parent(int index) {
    if (index > 0) {
        return (index - 1) / ARITY;
    }

    return -1;
}

__device__ uint64_t calcMerkle(const CuVector<BitArray>& leaves)
{
    LeaveHash l;
    PathHash p;

    int max_leaves = checked_next_power_of_n(leaves.size(), ARITY);

    int num_nodes = (max_leaves - 1) / (ARITY - 1);
    int tree_size = max_leaves + num_nodes;
    int tree_depth = calc_tree_depth(tree_size);
    int padding_depth = DEPTH - tree_depth;

    BitArray empty_hash = p.hash_empty();
    //print_bool("empty_hash:",empty_hash);

    // 计算树的最小尺寸，模拟原Rust代码逻辑
    int arity = ARITY;
    bool all_nodes_are_full = leaves.size() % arity == 0;
    int minimum_tree_size = num_nodes + leaves.size() + (all_nodes_are_full ? 0 : arity - leaves.size() % arity);
    if (minimum_tree_size < 1)
        minimum_tree_size = 1;

    CuVector<BitArray> tree(minimum_tree_size, empty_hash);
    for (int i = 0; i < leaves.size(); i++)
    {
        int idx = i + num_nodes;
        tree[idx] = l.hash_leaf(leaves[i]);
    }

    int start_index = num_nodes;
    while (true) {
        int start = parent(start_index);
        if (start == -1) {
            break;
        }

        Range end_pair = child_indexes(start);
        int end = end_pair.r0;

        for (int node_idx = start; node_idx < end; ++node_idx)
        {
            CuVector<BitArray> child_nodes;
            Range child_index_range = child_indexes(node_idx);
            for (int child_idx = child_index_range.r0; child_idx < child_index_range.r1; child_idx++)
            {
                if (child_idx < tree.size())
                    child_nodes.push_back(tree[child_idx]);
                else
                    child_nodes.push_back(empty_hash);
            }
            tree[node_idx] = p.hash_children(child_nodes);
        }

        start_index = start;
    }

    BitArray current_root_hash = tree[0];
    for (int i = 0; i < padding_depth; ++i) {
        CuVector<BitArray> input(ARITY, empty_hash);
        input[0] = current_root_hash;
        current_root_hash = p.hash_children(input);
    }
    //dump_merkle(tree);

    BitArray root_hash = current_root_hash.mid(0, 64);
    uint64_t u = 0;
    for (int i = 0; i < root_hash.size(); i++)
    {
        u |= (uint64_t)root_hash[i].value << (63 - i);
    }
    return u;
}